/**
 *@file yolo5n6_plugin.cpp
 *@author lynxi
 *@version v1.0
 *@date 2021-05-18
 *@par Copyright:
 *© 2018 北京灵汐科技有限公司 版权所有。\n
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。\n
 *© 2018 Lynxi Technologies Co., Ltd. All rights reserved. \n
 * NOTICE: All information contained here is, and remains the property of Lynxi.
 *This file can not be copied or distributed without the permission of Lynxi
 *Technologies Co., Ltd.
 *@brief
 */
#include <sys/time.h>
#include <algorithm>
#include <cmath>
#include <future>
#include <iomanip>
#include <iostream>
#include <string>
#include <string.h>
#include <utility>
#include <vector>
#include "typeConv.hpp"
#include "common.h"
#include "yolox_post_process.h"
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

#define CLASS_NUM 80  // 2//80

using namespace COMMON;

struct Yolo5n6Config {
  std::vector<int> strides;
  std::vector<std::vector<std::pair<double, double>>> anchors_table;
  int class_num;
  std::vector<std::string> class_names;
};

static const Yolo5n6Config default_yolo5n6_config = {
    {8, 16, 32, 64},
    {{{19, 27}, {44, 40}, {38, 94}},
     {{96, 68}, {86, 152}, {180, 137}},
     {{140, 301}, {303, 264}, {238, 542}},
     {{436, 615}, {739, 380}, {925, 792}}},
    CLASS_NUM,
    {"person",        "bicycle",      "car",
     "motorcycle",    "airplane",     "bus",
     "train",         "truck",        "boat",
     "traffic light", "fire hydrant", "stop sign",
     "parking meter", "bench",        "bird",
     "cat",           "dog",          "horse",
     "sheep",         "cow",          "elephant",
     "bear",          "zebra",        "giraffe",
     "backpack",      "umbrella",     "handbag",
     "tie",           "suitcase",     "frisbee",
     "skis",          "snowboard",    "sports ball",
     "kite",          "baseball bat", "baseball glove",
     "skateboard",    "surfboard",    "tennis racket",
     "bottle",        "wine glass",   "cup",
     "fork",          "knife",        "spoon",
     "bowl",          "banana",       "apple",
     "sandwich",      "orange",       "broccoli",
     "carrot",        "hot dog",      "pizza",
     "donut",         "cake",         "chair",
     "couch",         "potted plant", "bed",
     "dining table",  "toilet",       "tv",
     "laptop",        "mouse",        "remote",
     "keyboard",      "cell phone",   "microwave",
     "oven",          "toaster",      "sink",
     "refrigerator",  "book",         "clock",
     "vase",          "scissors",     "teddy bear",
     "hair drier",    "toothbrush"}};

template <class ForwardIterator>
inline size_t argmin(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::min_element(first, last));
}

template <class ForwardIterator>
inline size_t argmax(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::max_element(first, last));
}

typedef struct Bbox {
  float xmin;
  float ymin;
  float xmax;
  float ymax;

  Bbox() {}

  Bbox(float xmin, float ymin, float xmax, float ymax)
      : xmin(xmin), ymin(ymin), xmax(xmax), ymax(ymax) {}

  ~Bbox() {}
} Bbox;

typedef struct Detection {
  int id;
  float score;
  Bbox bbox;
  const char *class_name;
  Detection() {}

  Detection(int id, float score, Bbox bbox)
      : id(id), score(score), bbox(bbox) {}

  Detection(int id, float score, Bbox bbox, const char *class_name)
      : id(id), score(score), bbox(bbox), class_name(class_name) {}

  friend bool operator>(const Detection &lhs, const Detection &rhs) {
    return (lhs.score > rhs.score);
  }

  ~Detection() {}
} Detection;

void yolo5n6Nms(std::vector<Detection> &input, float iou_threshold, int top_k,
                std::vector<Detection> &result, bool suppress) {
  std::stable_sort(input.begin(), input.end(), std::greater<Detection>());

  std::vector<bool> skip(input.size(), false);

  std::vector<float> areas;
  areas.reserve(input.size());
  for (size_t i = 0; i < input.size(); i++) {
    float width = input[i].bbox.xmax - input[i].bbox.xmin;
    float height = input[i].bbox.ymax - input[i].bbox.ymin;
    areas.push_back(width * height);
  }

  int count = 0;
  for (size_t i = 0; /*count < top_k && */ i < skip.size(); i++) {
    if (skip[i]) {
      continue;
    }
    skip[i] = true;
    ++count;

    for (size_t j = i + 1; j < skip.size(); ++j) {
      if (skip[j]) {
        continue;
      }
      if (suppress == false) {
        if (input[i].id != input[j].id) {
          continue;
        }
      }

      float xx1 = std::max(input[i].bbox.xmin, input[j].bbox.xmin);
      float yy1 = std::max(input[i].bbox.ymin, input[j].bbox.ymin);
      float xx2 = std::min(input[i].bbox.xmax, input[j].bbox.xmax);
      float yy2 = std::min(input[i].bbox.ymax, input[j].bbox.ymax);

      if (xx2 > xx1 && yy2 > yy1) {
        float area_intersection = (xx2 - xx1) * (yy2 - yy1);
        float iou_ratio =
            area_intersection / (areas[j] + areas[i] - area_intersection);
        if (iou_ratio > iou_threshold) {
          skip[j] = true;
        }
      }
    }
    result.push_back(input[i]);
  }
}

void tensorpostProcess(const Yolo5n6Config &yolo5n6_config, void *tensor, YoloxPostProcessInfo_t *post_info, int layer,
                       std::vector<Detection> &dets) {
  void *data = tensor;
  int stride = yolo5n6_config.strides[layer];
  int class_num = yolo5n6_config.class_num;
  int num_pred = class_num + 4 + 1;

  std::vector<int16_t> class_pred(yolo5n6_config.class_num, 0);

  std::vector<std::pair<double, double>> anchors =
      yolo5n6_config.anchors_table[layer];

  double fScaleW = (double)post_info->width / (double)post_info->ori_width;
  double fScaleH = (double)post_info->height / (double)post_info->ori_height;
  double fScale = std::min(fScaleW, fScaleH);

  int iResizeWidth = static_cast<int>(((double)post_info->ori_width * fScale)) & ~1;
  int iResizeHeight = static_cast<int>(((double)post_info->ori_height * fScale)) & ~1;

  if (iResizeWidth > 1920 || iResizeHeight > 1080) {
    fScaleW = (double)1920 / (double)iResizeWidth;
    fScaleH = (double)1080 / (double)iResizeHeight;
    fScale = fScale * std::min(fScaleW, fScaleH);
  }

  int grid_height, grid_width;
  grid_height = post_info->height / stride;
  grid_width = post_info->width / stride;

  int16_t box_score_threshold = float2half(post_info->score_threshold);
  for (int h = 0; h < grid_height; h++) {
    for (int w = 0; w < grid_width; w++) {
      for (int k = 0; k < (int)anchors.size(); k++) {
        int16_t *cur_data = (int16_t *)data + k * num_pred;

        int16_t objness = cur_data[4];
        if (objness < box_score_threshold /*post_info->score_threshold*/) {
          continue;
        }

        for (int index = 0; index < class_num; ++index) {
          class_pred[index] = (cur_data[5 + index]);
        }

        int16_t id = argmax(class_pred.begin(), class_pred.end());

        float confidence = half2float(objness) * half2float(class_pred[id]);

        if (confidence < post_info->score_threshold) {
          continue;
        }

        float center_x = half2float(cur_data[0]);
        float center_y = half2float(cur_data[1]);
        float scale_x = half2float(cur_data[2]);
        float scale_y = half2float(cur_data[3]);

        double xmin = (center_x - scale_x / 2.0);
        double ymin = (center_y - scale_y / 2.0);
        double xmax = (center_x + scale_x / 2.0);
        double ymax = (center_y + scale_y / 2.0);

        double xmin_org = xmin / fScale;
        double xmax_org = xmax / fScale;
        double ymin_org = ymin / fScale;
        double ymax_org = ymax / fScale;

        if (xmax_org <= 0 || ymax_org <= 0) {
          continue;
        }

        if (xmin_org > xmax_org || ymin_org > ymax_org) {
          continue;
        }

        xmin_org = std::max(xmin_org, 0.0);
        xmax_org = std::min(xmax_org, post_info->ori_width - 1.0);
        ymin_org = std::max(ymin_org, 0.0);
        ymax_org = std::min(ymax_org, post_info->ori_height - 1.0);

        Bbox bbox(xmin_org, ymin_org, xmax_org, ymax_org);
        dets.push_back(
            Detection((int)id, confidence, bbox,
                      yolo5n6_config.class_names[(int)id].c_str()));
      }
      data = (int16_t *)data + num_pred * anchors.size();
    }
  }
}

int lynYoloxPostProcess(YoloxPostProcessInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;

#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
#endif

  Yolo5n6Config yolo5n6_config = default_yolo5n6_config;
  generateClassNameAndNum(post_info->labelList,
                          yolo5n6_config.class_names,
                          yolo5n6_config.class_num);

  int num_pred = yolo5n6_config.class_num + 5;

  void *tensor0 = NULL;
  tensor0 = output_tensor;
  void *tensor1 = NULL;
  int grid_width1 = post_info->width / yolo5n6_config.strides[0];
  int grid_height1 = post_info->height / yolo5n6_config.strides[0];
  tensor1 =
      (int16_t *)tensor0 + 1 * 3 * grid_width1 * grid_height1 * num_pred;
  void *tensor2 = NULL;
  int grid_width2 = post_info->width / yolo5n6_config.strides[1];
  int grid_height2 = post_info->height / yolo5n6_config.strides[1];
  tensor2 =
      (int16_t *)tensor1 + 1 * 3 * grid_width2 * grid_height2 * num_pred;
  void *tensor3 = NULL;
  int grid_width3 = post_info->width / yolo5n6_config.strides[2];
  int grid_height3 = post_info->height / yolo5n6_config.strides[2];
  tensor3 =
      (int16_t *)tensor2 + 1 * 3 * grid_width3 * grid_height3 * num_pred;
  std::vector<Detection> dets;
  std::vector<Detection> det_restuls;
  tensorpostProcess(yolo5n6_config, tensor0, post_info, 0, dets);
  tensorpostProcess(yolo5n6_config, tensor1, post_info, 1, dets);
  tensorpostProcess(yolo5n6_config, tensor2, post_info, 2, dets);
  tensorpostProcess(yolo5n6_config, tensor3, post_info, 3, dets);

  yolo5n6Nms(dets, post_info->nms_threshold, post_info->nms_top_k, det_restuls,
             false);
  boxesInfo->boxesNum = std::min((int)det_restuls.size(), BOX_MAX_NUM);

  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    boxesInfo->boxes[i].xmax = det_restuls[i].bbox.xmax;
    boxesInfo->boxes[i].xmin = det_restuls[i].bbox.xmin;
    boxesInfo->boxes[i].ymax = det_restuls[i].bbox.ymax;
    boxesInfo->boxes[i].ymin = det_restuls[i].bbox.ymin;
    boxesInfo->boxes[i].score = det_restuls[i].score;
    boxesInfo->boxes[i].id = det_restuls[i].id;
    strncpy(boxesInfo->boxes[i].label, det_restuls[i].class_name,
            sizeof(boxesInfo->boxes[i].label) - 1);
    if (static_cast<int>(boxesInfo->boxes[i].xmax) > post_info->ori_width - 1)
      boxesInfo->boxes[i].xmax = post_info->ori_width - 1;
    if (static_cast<int>(boxesInfo->boxes[i].xmin) < 0)
      boxesInfo->boxes[i].xmin = 0;
    if (static_cast<int>(boxesInfo->boxes[i].ymax) > post_info->ori_height - 1)
      boxesInfo->boxes[i].ymax = post_info->ori_height - 1;
    if (static_cast<int>(boxesInfo->boxes[i].ymin) < 0)
      boxesInfo->boxes[i].ymin = 0;
  }

  return 0;
}
